#ifndef QCOW2_H
#define QCOW2_H

#include <stdio.h>
#include <stdint.h>
#include <list>
#include <vector>
#include <map>
#include <string>
#include "common.h"
#include "block.h"

#define QCOW_MAGIC (('Q' << 24) | ('F' << 16) | ('I' << 8) | 0xfb)
#define QCOW_VERSION 2
#define QCOW_VERSION3 3

#define QCOW_CRYPT_NONE 0
#define QCOW_CRYPT_AES  1

#define QCOW_OFLAG_COMPRESSED (1ULL << 63)
#define QCOW_OFLAG_COPIED     (1ULL << 62)

#define L1E_OFFSET_MASK 0x00ffffffffffff00ULL

// 簇类型定义
#define QCOW2_CLUSTER_UNALLOCATED 0
#define QCOW2_CLUSTER_NORMAL 1
#define QCOW2_CLUSTER_COMPRESSED 2
#define QCOW2_CLUSTER_ZERO 3

// 引用计数表相关
#define REFCOUNT_TABLE_ENTRIES_PER_CLUSTER(size) ((size) / sizeof(uint64_t))
#define REFCOUNT_TABLE_OFFSET_MASK 0xFFFFFFFFFFFF0000ULL

// 错误码
#define QCOW2_ERROR_IO -1
#define QCOW2_ERROR_INVALID -2
#define QCOW2_ERROR_NOSPC -3
#define QCOW2_ERROR_CORRUPT -4
#define QCOW2_ERROR_L1_FULL -5
#define QCOW2_ERROR_EXIST -6

// QCow2 Version 3 头文件
typedef struct QCowHeaderV3 {
    uint32_t magic;
    uint32_t version;
    uint64_t backing_file_offset;
    uint32_t backing_file_size;
    uint32_t cluster_bits;
    uint64_t size;
    uint32_t crypt_method;
    uint32_t l1_size;
    uint64_t l1_table_offset;
    uint64_t refcount_table_offset;
    uint32_t refcount_table_clusters;
    uint32_t nb_snapshots;
    uint64_t snapshots_offset;

    // Version 3 新增字段
    uint64_t incompatible_features;
    uint64_t compatible_features;
    uint64_t autoclear_features;
    uint32_t refcount_order;
    uint32_t header_length;
} QEMU_PACKED QCowHeaderV3;

// 引用计数块
typedef struct QCow2RefcountBlock {
    uint16_t *counts;           // 引用计数数组
    uint64_t offset;            // 在文件中的偏移
    int size;                   // 块大小
    bool dirty;                 // 脏标志
} QCow2RefcountBlock;

// 数据块信息结构
typedef struct DataBlockInfo {
    uint64_t offset;            // 在虚拟磁盘中的偏移
    uint64_t length;            // 数据块长度
    uint64_t file_offset;       // 在文件中的物理偏移
    int cluster_type;           // 簇类型
    bool allocated;             // 是否已分配
} DataBlockInfo;

typedef struct QCow2State {
    int cluster_bits;
    int cluster_size;
    int cluster_sectors;
    int l2_bits;
    int l2_size;
    int l1_size;
    int refcount_order;
    int crypt_method;
    int version;

    uint64_t cluster_offset_mask;

    // Version 3 特性
    uint64_t incompatible_features;
    uint64_t compatible_features;
    uint64_t autoclear_features;

    // 表结构
    uint64_t *l1_table;
    uint64_t l1_table_offset;

    uint64_t *refcount_table;
    uint64_t refcount_table_offset;
    int refcount_table_size;

    // 缓存
    uint64_t *l2_cache;
    uint64_t l2_cache_offsets[64];
    bool l2_cache_valid[64];
    int l2_cache_size;
    int l2_cache_index;

    std::map<uint64_t, QCow2RefcountBlock*> refcount_block_cache;

    uint8_t *cluster_cache;
    uint64_t cluster_cache_offset;

    BlockDriverState *bs;
} QCow2State;

#ifdef __cplusplus
extern "C" {
#endif

// 基本操作
int qcow2_open(BlockDriverState *bs, const char *filename);
int qcow2_close(BlockDriverState *bs);
int64_t qcow2_getlength(BlockDriverState *bs);
int qcow2_read(BlockDriverState *bs, int64_t offset, uint8_t *buf, int64_t bytes);
int qcow2_probe(const uint8_t *buf, int buf_size);
int qcow2_create(const char *filename, int64_t size, const char *backing_file,
                int cluster_size, int version);

// 写操作
int qcow2_write(BlockDriverState *bs, int64_t offset, const uint8_t *buf, int64_t bytes);
int qcow2_allocate_clusters(BlockDriverState *bs, uint64_t offset, uint64_t size);
int qcow2_pwritev(BlockDriverState *bs, uint64_t offset, uint64_t bytes,
                  const uint8_t *buf, uint64_t flags);

// 簇管理
uint64_t qcow2_alloc_clusters(BlockDriverState *bs, uint64_t size);
int qcow2_free_clusters(BlockDriverState *bs, uint64_t offset, uint64_t size);
int qcow2_update_cluster_refcount(BlockDriverState *bs, uint64_t cluster_offset, int add);

// 表管理
int qcow2_grow_l1_table(BlockDriverState *bs, int min_size);
int qcow2_grow_refcount_table(BlockDriverState *bs, int min_size);
int qcow2_alloc_cluster_at(BlockDriverState *bs, uint64_t offset, uint64_t *cluster_offset);

// 引用计数
int qcow2_update_refcount(BlockDriverState *bs, uint64_t offset, uint64_t length, int add);
int qcow2_get_refcount(BlockDriverState *bs, uint64_t cluster_offset);
int qcow2_alloc_refcount_block(BlockDriverState *bs, uint64_t cluster_offset);

// 数据块查询
std::list<DataBlockInfo*> qcow2_query_data_blocks(BlockDriverState *bs);
std::list<DataBlockInfo*> qcow2_query_data_blocks_fast(BlockDriverState *bs);

// 工具函数
int qcow2_flush(BlockDriverState *bs);
int qcow2_check_consistency(BlockDriverState *bs);
uint64_t qcow2_virtual_size(BlockDriverState *bs);
uint64_t qcow2_physical_size(BlockDriverState *bs);

// backing file 管理
int qcow2_set_backing_file(BlockDriverState *bs, const char *backing_file);
char* qcow2_get_backing_file(BlockDriverState *bs);

// 内部函数
int qcow2_read_l1_table(QCow2State *s, FILE *f);
int qcow2_read_l2_table(QCow2State *s, FILE *f, uint64_t l2_offset, uint64_t **l2_table);
int qcow2_read_refcount_table(QCow2State *s, FILE *f);
int qcow2_write_l1_table(QCow2State *s);
int qcow2_write_refcount_table(QCow2State *s);
int qcow2_update_l2_entry(BlockDriverState *bs, uint64_t l2_offset, int l2_index,
                         uint64_t cluster_offset);

#ifdef __cplusplus
}
#endif

#endif